import torch
import torchvision.datasets
from torch.nn import ReLU, Sigmoid, SELU
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from torchvision.datasets import CIFAR10
from torchvision.transforms import ToTensor

input = torch.tensor([
    [1, -0.5, 3],
    [2, -4, 4],
    [5, -3, 4]
])

input = torch.reshape(input, (-1, 1, 3, 3))
"""
非线性变换，是为了让图片看起来更统一一些，训练出来的模型适用性更广泛一些
"""
relu = ReLU()
output = relu(input)
print(output)

class ReluModule(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.relu = ReLU()
        self.sigmoid = Sigmoid()
        self.selu = SELU()

    def forward(self, input, type):
        if type == 1:
            return self.relu(input)
        if type == 2:
            return self.sigmoid(input)
        if type == 3:
            return self.selu(input)

reluModule = ReluModule()
dataset = CIFAR10(root="datasets", train=False, transform=ToTensor(), download=True)
dataloader = DataLoader(dataset=dataset, batch_size=64, drop_last=False)

writer = SummaryWriter("logs")
step = 0

for imgs, targets in dataloader:
    writer.add_images("input", imgs, step)
    writer.add_images("none_liner_relu", reluModule(imgs, 1), step)
    writer.add_images("none_liner_sigmoid", reluModule(imgs, 2), step)
    writer.add_images("none_liner_selu", reluModule(imgs, 3), step)
    step += 1

writer.close()
print("success")

